# 邓迪心第一版数据集类
from typing import Tuple
import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor
from apps.dhlp.dxd_data_v1 import DxdDataV1

class Ddxv1Ds(Dataset):
    def __init__(self):
        super(Ddxv1Ds, self).__init__()
        # 读入work/datasets/dhlp/v1/ds.txt文件内容
        ds_fn = 'work/datasets/dhlp/v1/ds.txt'
        self.recs = []
        with open(ds_fn, 'r', encoding='utf-8') as rfd:
            for row in rfd:
                row = row.strip()
                arrs0 = row.split(',')
                self.recs.append({
                    'fn': arrs0[DxdDataV1.DSF_FN],
                    'did': arrs0[DxdDataV1.DSF_DID],
                    'rid': arrs0[DxdDataV1.DSF_RID],
                    'label': arrs0[DxdDataV1.DSF_LABEL],
                    'pos': arrs0[DxdDataV1.DSF_POS],
                    'idx': arrs0[DxdDataV1.DSF_IDX],
                    'cls': arrs0[DxdDataV1.DSF_CLS]
                })

    def __len__(self) -> int:
        return len(self.recs)

    def __getitem__(self, idx:int) -> Tuple[torch.Tensor, torch.Tensor]:
        X = torch.load(self.recs[idx]['fn'], weights_only=True)
        y = int(self.recs[idx]['cls'])
        return X, y